import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
import torch
import time
from transformers import AutoModel, AutoProcessor
from transformers.image_utils import load_image
from PIL import Image
import threading
from queue import Queue, Empty
import numpy as np
from tqdm import tqdm
import h5py
import logging
from pathlib import Path

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('embedding_extraction.log'),
        logging.StreamHandler()
    ]
)

# 配置
NUM_THREADS = 10
BATCH_SIZE = 32  # 批处理大小

class EmbeddingExtractor:
    def __init__(self, model_name, num_threads=NUM_THREADS):
        self.device = torch.device("cuda")
        self.model = AutoModel.from_pretrained(model_name, device_map="cuda").vision_model
        self.processor = AutoProcessor.from_pretrained(model_name)
        self.num_threads = num_threads
        self.queue = Queue()
        self.results = {}
        self.lock = threading.Lock()
        self.failed_images = []
        self.pbar = None
        
    def process_image(self, image_path):
        """处理单个图片"""
        try:
            image = Image.open(image_path).convert('RGB')
            inputs = self.processor(images=[image], return_tensors="pt").to(self.device)
            with torch.no_grad():
                embeddings = self.model(**inputs).pooler_output
            return embeddings.cpu().numpy()
        except Exception as e:
            logging.error(f"Error processing {image_path}: {str(e)}")
            return None

    def process_batch(self, image_paths):
        """批量处理图片"""
        try:
            images = [Image.open(path).convert('RGB') for path in image_paths]
            inputs = self.processor(images=images, return_tensors="pt").to(self.device)
            with torch.no_grad():
                embeddings = self.model(**inputs).pooler_output
            return embeddings.cpu().numpy()
        except Exception as e:
            logging.error(f"Error processing batch: {str(e)}")
            return None

    def worker(self):
        """工作线程函数"""
        local_batch = []
        while True:
            try:
                # 获取下一个任务
                item = self.queue.get(timeout=5)
                local_batch.append(item)
                
                # 如果批次满了或队列空了，处理当前批次
                if len(local_batch) >= BATCH_SIZE or self.queue.empty():
                    if local_batch:
                        embeddings = self.process_batch([item[1] for item in local_batch])
                        if embeddings is not None:
                            with self.lock:
                                for (img_id, _), emb in zip(local_batch, embeddings):
                                    self.results[img_id] = emb
                                self.pbar.update(len(local_batch))
                        else:
                            self.failed_images.extend([item[0] for item in local_batch])
                    local_batch = []
                
                self.queue.task_done()
                
            except Empty:
                # 处理最后的批次
                if local_batch:
                    embeddings = self.process_batch([item[1] for item in local_batch])
                    if embeddings is not None:
                        with self.lock:
                            for (img_id, _), emb in zip(local_batch, embeddings):
                                self.results[img_id] = emb
                            self.pbar.update(len(local_batch))
                    else:
                        self.failed_images.extend([item[0] for item in local_batch])
                break
            except Exception as e:
                logging.error(f"Worker error: {str(e)}")
                if local_batch:
                    self.failed_images.extend([item[0] for item in local_batch])
                continue

    def extract_embeddings(self, image_dir, output_file):
        """提取目录下所有图片的embeddings"""
        try:
            # 获取所有图片
            image_files = list(Path(image_dir).glob('*.jpg'))
            total_images = len(image_files)
            logging.info(f"Found {total_images} images in {image_dir}")
            
            # 初始化进度条
            self.pbar = tqdm(total=total_images, desc=f"Processing {Path(image_dir).name}")
            
            # 填充队列
            for img_path in image_files:
                self.queue.put((img_path.stem, str(img_path)))
            
            # 创建线程
            threads = []
            for _ in range(self.num_threads):
                thread = threading.Thread(target=self.worker)
                thread.start()
                threads.append(thread)
            
            # 等待所有线程完成
            for thread in threads:
                thread.join()
            
            # 保存结果
            with h5py.File(output_file, 'w') as f:
                for img_id, embedding in self.results.items():
                    f.create_dataset(img_id, data=embedding)
            
            # 保存失败的图片信息
            if self.failed_images:
                with open(f"{output_file}_failed.txt", 'w') as f:
                    for img_id in self.failed_images:
                        f.write(f"{img_id}\n")
            
            logging.info(f"Completed! Processed {len(self.results)} images, failed {len(self.failed_images)} images")
            logging.info(f"Results saved to {output_file}")
            
        except Exception as e:
            logging.error(f"Error in extraction process: {str(e)}")
        finally:
            if self.pbar:
                self.pbar.close()

def main():
    # 配置
    model_name = "google/siglip2-so400m-patch14-384"
    images_path_q = '/home/tfj/datasets/image_retri10k/images/'
    images_path_v = '/home/tfj/datasets/image_retri10k/images_sup/'
    # images_path_o = '/home/tfj/datasets/image_retri10k/eval_images_o/'
    
    # 创建提取器实例
    extractor = EmbeddingExtractor(model_name, NUM_THREADS)
    
    # 处理Q集图片
    extractor.extract_embeddings(images_path_q, 'embeddings.h5')
    
    # 重置结果存储
    extractor.results = {}
    extractor.failed_images = []

    # 处理V集图片
    extractor.extract_embeddings(images_path_v, 'embeddings_sup.h5')

    # 重置结果存储
    # extractor.results = {}
    # extractor.failed_images = []

    # # 处理O集图片
    # extractor.extract_embeddings(images_path_o, 'embeddings_o.h5')
if __name__ == "__main__":
    main()
